{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f5ac3b08",
   "metadata": {},
   "source": [
    "Author:Mohammadreza Tavasoli Naeini Advisor:Hanrui Wang"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e1c8036",
   "metadata": {},
   "source": [
    "This tutorial shows; how to use the TrochQuantum library in building a quantum transformer.\n",
    "This notebook is adapted, excerpted from https://github.com/rdisipio/qtransformer, written by Riccardo Di Sipio. \n",
    "In this notebook, we use the TorchQuantum library instead of the Pennylane library to build a quantum machine learning model.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dc0dc60",
   "metadata": {},
   "source": [
    "## Introduction to Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7022d464",
   "metadata": {},
   "source": [
    "Nowadays, Transformer architecture is dominant in the field of sequential data processing, such as natural language processing. Transformers designed to effieciently analyze sequential data on GPUs. Therfore, we see stonishing results in translation, generation from the models build on transformer blocks.[1],[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98c6c7f9",
   "metadata": {},
   "source": [
    "Transformer block does not have memory to sequentially anaylize the data; however, it has position-dependent embedding, which is a kind of word embbeding that embed position of the input in the vector too. Furthermore, trnasformer block use self attation mechanisem to understand which part of the output is related to the input data. The core of transformer block is multi head attetion layer. In multi head atteation layer, three diffrent linear transformer, WQ,WK,and WV, applies to each word embbeding to transform them to new  internal representation called Query (Q), Key (K), and Value (V ). These states are then passed to the function that calculates the attention weights.[1],[2]\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c17fc45b",
   "metadata": {},
   "source": [
    "<figure>\n",
    "<img src=\"transformer.png\" style=\"width:30%\">\n",
    "<figcaption align = \"center\"> transformer architecture </figcaption>\n",
    "</figure>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7caf96fb",
   "metadata": {},
   "source": [
    "## Quantum Transformer\n",
    "\n",
    "We convert the classical transformer to the quantum enhanced transformer by replacing the linear transformations WQ, WK, and WV with Varitinal Quantum Circuits. [1] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0797838",
   "metadata": {},
   "source": [
    "## Build and train a Quantum Transformer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4666fb63",
   "metadata": {},
   "source": [
    "## Installation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da4d600b",
   "metadata": {},
   "source": [
    "Install torchquantum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "38e1421a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'torchquantum'...\n",
      "remote: Enumerating objects: 11649, done.\u001b[K\n",
      "remote: Counting objects: 100% (539/539), done.\u001b[K\n",
      "remote: Compressing objects: 100% (273/273), done.\u001b[K\n",
      "remote: Total 11649 (delta 315), reused 456 (delta 255), pack-reused 11110\u001b[K\n",
      "Receiving objects: 100% (11649/11649), 6.43 MiB | 3.37 MiB/s, done.\n",
      "Resolving deltas: 100% (6472/6472), done.\n"
     ]
    }
   ],
   "source": [
    "!git clone https://github.com/mit-han-lab/torchquantum.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "423f8bae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/mohammad/.Trash/torchquantum 10-11-35-421/torchquantum\n"
     ]
    }
   ],
   "source": [
    "%cd torchquantum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "673fd538",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Obtaining file:///Users/mohammad/.Trash/torchquantum%2010-11-35-421/torchquantum\n",
      "Requirement already satisfied: numpy>=1.19.2 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (1.20.1)\n",
      "Requirement already satisfied: torchvision>=0.9.0.dev20210130 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (0.9.1)\n",
      "Requirement already satisfied: tqdm>=4.56.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (4.60.0)\n",
      "Requirement already satisfied: setuptools>=52.0.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (65.3.0)\n",
      "Requirement already satisfied: torch>=1.8.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (1.8.1)\n",
      "Requirement already satisfied: torchpack>=0.3.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (0.3.1)\n",
      "Requirement already satisfied: qiskit>=0.32.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (0.38.0)\n",
      "Requirement already satisfied: matplotlib>=3.3.2 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (3.4.1)\n",
      "Requirement already satisfied: pathos>=0.2.7 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (0.2.9)\n",
      "Requirement already satisfied: pylatexenc>=2.10 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchquantum==0.1.2) (2.10)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from matplotlib>=3.3.2->torchquantum==0.1.2) (1.3.1)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from matplotlib>=3.3.2->torchquantum==0.1.2) (8.2.0)\n",
      "Requirement already satisfied: pyparsing>=2.2.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from matplotlib>=3.3.2->torchquantum==0.1.2) (2.4.7)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from matplotlib>=3.3.2->torchquantum==0.1.2) (2.8.1)\n",
      "Requirement already satisfied: cycler>=0.10 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from matplotlib>=3.3.2->torchquantum==0.1.2) (0.10.0)\n",
      "Requirement already satisfied: six in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from cycler>=0.10->matplotlib>=3.3.2->torchquantum==0.1.2) (1.15.0)\n",
      "Requirement already satisfied: ppft>=1.7.6.5 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from pathos>=0.2.7->torchquantum==0.1.2) (1.7.6.5)\n",
      "Requirement already satisfied: dill>=0.3.5.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from pathos>=0.2.7->torchquantum==0.1.2) (0.3.5.1)\n",
      "Requirement already satisfied: multiprocess>=0.70.13 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from pathos>=0.2.7->torchquantum==0.1.2) (0.70.13)\n",
      "Requirement already satisfied: pox>=0.3.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from pathos>=0.2.7->torchquantum==0.1.2) (0.3.1)\n",
      "Requirement already satisfied: qiskit-aer==0.11.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit>=0.32.0->torchquantum==0.1.2) (0.11.0)\n",
      "Requirement already satisfied: qiskit-terra==0.21.2 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit>=0.32.0->torchquantum==0.1.2) (0.21.2)\n",
      "Requirement already satisfied: qiskit-ibmq-provider==0.19.2 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit>=0.32.0->torchquantum==0.1.2) (0.19.2)\n",
      "Requirement already satisfied: scipy>=1.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-aer==0.11.0->qiskit>=0.32.0->torchquantum==0.1.2) (1.6.1)\n",
      "Requirement already satisfied: urllib3>=1.21.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (1.26.4)\n",
      "Requirement already satisfied: requests>=2.19 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (2.25.1)\n",
      "Requirement already satisfied: websockets>=10.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (10.3)\n",
      "Requirement already satisfied: websocket-client>=1.0.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (1.4.1)\n",
      "Requirement already satisfied: requests-ntlm>=1.1.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (1.1.0)\n",
      "Requirement already satisfied: psutil>=5 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-terra==0.21.2->qiskit>=0.32.0->torchquantum==0.1.2) (5.8.0)\n",
      "Requirement already satisfied: sympy>=1.3 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-terra==0.21.2->qiskit>=0.32.0->torchquantum==0.1.2) (1.7.1)\n",
      "Requirement already satisfied: retworkx>=0.11.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-terra==0.21.2->qiskit>=0.32.0->torchquantum==0.1.2) (0.11.0)\n",
      "Requirement already satisfied: tweedledum<2.0,>=1.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-terra==0.21.2->qiskit>=0.32.0->torchquantum==0.1.2) (1.1.1)\n",
      "Requirement already satisfied: stevedore>=3.0.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-terra==0.21.2->qiskit>=0.32.0->torchquantum==0.1.2) (4.0.0)\n",
      "Requirement already satisfied: symengine>=0.9 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-terra==0.21.2->qiskit>=0.32.0->torchquantum==0.1.2) (0.9.2)\n",
      "Requirement already satisfied: ply>=3.10 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from qiskit-terra==0.21.2->qiskit>=0.32.0->torchquantum==0.1.2) (3.11)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from requests>=2.19->qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (2020.12.5)\n",
      "Requirement already satisfied: chardet<5,>=3.0.2 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from requests>=2.19->qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (4.0.0)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from requests>=2.19->qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (2.10)\n",
      "Requirement already satisfied: ntlm-auth>=1.0.2 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from requests-ntlm>=1.1.0->qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (1.5.0)\n",
      "Requirement already satisfied: cryptography>=1.3 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from requests-ntlm>=1.1.0->qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (3.4.7)\n",
      "Requirement already satisfied: cffi>=1.12 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from cryptography>=1.3->requests-ntlm>=1.1.0->qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (1.14.5)\n",
      "Requirement already satisfied: pycparser in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from cffi>=1.12->cryptography>=1.3->requests-ntlm>=1.1.0->qiskit-ibmq-provider==0.19.2->qiskit>=0.32.0->torchquantum==0.1.2) (2.20)\n",
      "Requirement already satisfied: pbr!=2.1.0,>=2.0.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from stevedore>=3.0.0->qiskit-terra==0.21.2->qiskit>=0.32.0->torchquantum==0.1.2) (5.10.0)\n",
      "Requirement already satisfied: mpmath>=0.19 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from sympy>=1.3->qiskit-terra==0.21.2->qiskit>=0.32.0->torchquantum==0.1.2) (1.2.1)\n",
      "Requirement already satisfied: typing-extensions in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torch>=1.8.0->torchquantum==0.1.2) (3.7.4.3)\n",
      "Requirement already satisfied: multimethod in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchpack>=0.3.0->torchquantum==0.1.2) (1.8)\n",
      "Requirement already satisfied: tensorboard in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchpack>=0.3.0->torchquantum==0.1.2) (2.9.1)\n",
      "Requirement already satisfied: tensorpack in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchpack>=0.3.0->torchquantum==0.1.2) (0.11)\n",
      "Requirement already satisfied: h5py in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchpack>=0.3.0->torchquantum==0.1.2) (3.1.0)\n",
      "Requirement already satisfied: pyyaml in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchpack>=0.3.0->torchquantum==0.1.2) (5.3.1)\n",
      "Requirement already satisfied: toml in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchpack>=0.3.0->torchquantum==0.1.2) (0.10.2)\n",
      "Requirement already satisfied: loguru in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from torchpack>=0.3.0->torchquantum==0.1.2) (0.6.0)\n",
      "Requirement already satisfied: wheel>=0.26 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (0.36.2)\n",
      "Requirement already satisfied: grpcio>=1.24.3 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (1.37.0)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (1.28.0)\n",
      "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (1.8.0)\n",
      "Requirement already satisfied: protobuf<3.20,>=3.9.2 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (3.15.7)\n",
      "Requirement already satisfied: absl-py>=0.4 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (1.2.0)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (3.3.4)\n",
      "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (0.6.1)\n",
      "Requirement already satisfied: werkzeug>=1.0.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (1.0.1)\n",
      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (0.4.4)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from google-auth<3,>=1.6.3->tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (4.7.2)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from google-auth<3,>=1.6.3->tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (0.2.8)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from google-auth<3,>=1.6.3->tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (4.2.1)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (1.3.0)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (0.4.8)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard->torchpack>=0.3.0->torchquantum==0.1.2) (3.1.0)\n",
      "Requirement already satisfied: pyzmq>=16 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorpack->torchpack>=0.3.0->torchquantum==0.1.2) (22.3.0)\n",
      "Requirement already satisfied: tabulate>=0.7.7 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorpack->torchpack>=0.3.0->torchquantum==0.1.2) (0.8.10)\n",
      "Requirement already satisfied: msgpack>=0.5.2 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorpack->torchpack>=0.3.0->torchquantum==0.1.2) (1.0.4)\n",
      "Requirement already satisfied: termcolor>=1.1 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorpack->torchpack>=0.3.0->torchquantum==0.1.2) (1.1.0)\n",
      "Requirement already satisfied: msgpack-numpy>=0.4.4.2 in /Users/mohammad/miniforge3/lib/python3.8/site-packages (from tensorpack->torchpack>=0.3.0->torchquantum==0.1.2) (0.4.8)\n",
      "Installing collected packages: torchquantum\n",
      "  Attempting uninstall: torchquantum\n",
      "    Found existing installation: torchquantum 0.1.2\n",
      "    Uninstalling torchquantum-0.1.2:\n",
      "      Successfully uninstalled torchquantum-0.1.2\n",
      "  Running setup.py develop for torchquantum\n",
      "Successfully installed torchquantum\n"
     ]
    }
   ],
   "source": [
    "!pip install --editable ."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e24ab702",
   "metadata": {},
   "source": [
    "## Setup\n",
    "Our code requires torchquantum lib, and pytorch. We need torch, function from torch.nn.functional, optimizers(optim), torchquantum module, and torchtext."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04ec45a0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3e6c7e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "import torchquantum as tq\n",
    "import torchquantum.functional as tqf\n",
    "import argparse\n",
    "import tqdm\n",
    "import time\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torchtext.legacy\n",
    "from torchtext.legacy import data, datasets, vocab"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a25a2db2",
   "metadata": {},
   "source": [
    "## Build MultiHead Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cf2bc730",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttentionBase(nn.Module):\n",
    "    def __init__(self,\n",
    "                 embed_dim: int,\n",
    "                 num_heads: int,\n",
    "                 dropout: float = 0.1,\n",
    "                 mask=None,\n",
    "                 use_bias=False):\n",
    "        super(MultiHeadAttentionBase, self).__init__()\n",
    "\n",
    "        assert embed_dim % num_heads == 0, f\"Embedding dimension ({embed_dim}) should be divisible by number of heads ({num_heads})\"\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.d_k = embed_dim // num_heads  # projection dimensions\n",
    "        self.k_linear = None\n",
    "        self.q_linear = None\n",
    "        self.v_linear = None\n",
    "        self.combine_heads = None\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.attn_weights = None\n",
    "    \n",
    "    def separate_heads(self, x):\n",
    "        '''\n",
    "        split into N heads\n",
    "        from (batch_size, seq_len, embed_dim)\n",
    "        to   (batch_size, seq_len, num_heads, embed_dim)\n",
    "        then transpose (1,2) to (batch_size, num_heads, seq_len, embed_dim)\n",
    "        to make mat mult straightforward for each head\n",
    "        '''\n",
    "        batch_size = x.size(0)\n",
    "        x = x.view(batch_size, -1, self.num_heads, self.d_k)\n",
    "        return x.transpose(1, 2)\n",
    "\n",
    "    def attention(self, query, key, value, mask=None, dropout=None):\n",
    "        '''\n",
    "        Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k))V\n",
    "        '''\n",
    "        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.d_k)\n",
    "        # see also: https://tensorchiefs.github.io/dlday2018/tutorial/einsum.html\n",
    "        #scores = torch.einsum('bijh, bkjh -> bikh', query, key) / math.sqrt(self.d_k)\n",
    "        if mask is not None:\n",
    "            mask = mask.unsqueeze(1)\n",
    "            scores = scores.masked_fill(mask == 0, -1e9)\n",
    "        scores = F.softmax(scores, dim=-1)\n",
    "        if dropout is not None:\n",
    "            scores = dropout(scores)\n",
    "        attn = torch.matmul(scores, value)\n",
    "        return attn, scores\n",
    "    \n",
    "    def downstream(self, query, key, value, batch_size, mask=None):\n",
    "        Q = self.separate_heads(query)\n",
    "        K = self.separate_heads(key)\n",
    "        V = self.separate_heads(value)\n",
    "\n",
    "        x, self.attn_weights = self.attention(Q, K, V, mask, dropout=self.dropout)\n",
    "\n",
    "        concat = x.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_dim)\n",
    "\n",
    "        return concat\n",
    "        # output = self.combine_heads(concat)\n",
    "        # return output\n",
    "\n",
    "   # def forward(self, x, mask=None):\n",
    "    #    raise NotImplementedError(\"Base class does not execute forward function.\")\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c1bc0e4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttentionClassical(MultiHeadAttentionBase):\n",
    "    \n",
    "    def __init__(self, embed_dim: int,\n",
    "                 num_heads: int,\n",
    "                 dropout=0.1,\n",
    "                 mask=None,\n",
    "                 use_bias=False):\n",
    "        super(MultiHeadAttentionClassical, self).__init__(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout, mask=mask, use_bias=use_bias)\n",
    "\n",
    "        self.k_linear = nn.Linear(embed_dim, embed_dim, bias=use_bias)\n",
    "        self.q_linear = nn.Linear(embed_dim, embed_dim, bias=use_bias)\n",
    "        self.v_linear = nn.Linear(embed_dim, embed_dim, bias=use_bias)\n",
    "        self.combine_heads = nn.Linear(embed_dim, embed_dim, bias=use_bias)\n",
    "    \n",
    "    def forward(self, x, mask=None):\n",
    "        batch_size, seq_len, embed_dim = x.size()\n",
    "        assert embed_dim == self.embed_dim, f\"Input embedding ({embed_dim}) does not match layer embedding size ({self.embed_dim})\"\n",
    "\n",
    "        K = self.k_linear(x)\n",
    "        Q = self.q_linear(x)\n",
    "        V = self.v_linear(x)\n",
    "\n",
    "        x = self.downstream(Q, K, V, batch_size, mask)\n",
    "        output = self.combine_heads(x)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9decfa8c",
   "metadata": {},
   "source": [
    "Typically, a Quantum Neural Network module consists of three parts: encoder, ansatz, and measurement. We can create an encoder by passing a list of gates to tq.GeneralEncoder. Each entry in the list contains input_idx, func, and wires. Here, each qubit has a rotation-X gate. 8 RX gates in total. They can encode the 8 input data to the quantum state. Then we choose ansatz such that they are entangled between each other, rotated by an arbitrary angle. Finally, we perform Pauli-Z measurements on each qubit by creating a tq.MeasureAll module and passing tq.PauliZ to it. The measure function will return four expectation values from four qubits.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7ad3421c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttentionQuantum(MultiHeadAttentionBase):\n",
    "    class QLayer(tq.QuantumModule):\n",
    "        def __init__(self):\n",
    "            super().__init__()    \n",
    "            self.n_wires = 8\n",
    "            self.encoder = tq.GeneralEncoder(\n",
    "        [   {'input_idx': [0], 'func': 'rx', 'wires': [0]},\n",
    "            {'input_idx': [1], 'func': 'rx', 'wires': [1]},\n",
    "            {'input_idx': [2], 'func': 'rx', 'wires': [2]},\n",
    "            {'input_idx': [3], 'func': 'rx', 'wires': [2]},\n",
    "            {'input_idx': [4], 'func': 'rx', 'wires': [4]},\n",
    "            {'input_idx': [5], 'func': 'rx', 'wires': [5]},\n",
    "             {'input_idx': [6], 'func': 'rx', 'wires': [6]},\n",
    "          {'input_idx': [7], 'func': 'rx', 'wires': [7]},\n",
    "        ])\n",
    "            self.rx0 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx1 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx2 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx3 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx4 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx5 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx6 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx7 = tq.RX(has_params=True, trainable=True)\n",
    "            self.measure = tq.MeasureAll(tq.PauliZ)\n",
    "        def forward (self, x, q_device: tq.QuantumDevice):\n",
    "            self.q_device = q_device\n",
    "            self.encoder(self.q_device, x)\n",
    "            self.rx0(self.q_device, wires=0)\n",
    "            self.rx1(self.q_device, wires=1)\n",
    "            self.rx2(self.q_device, wires=2)\n",
    "            self.rx3(self.q_device, wires=3)\n",
    "            self.rx4(self.q_device, wires=4)\n",
    "            self.rx5(self.q_device, wires=5)\n",
    "            self.rx6(self.q_device, wires=6)\n",
    "            self.rx7(self.q_device, wires=7)\n",
    "            for k in range(self.n_wires):\n",
    "                if k==self.n_wires-1:\n",
    "                    tqf.cnot(self.q_device, wires=[k, 0]) \n",
    "                else:\n",
    "                    tqf.cnot(self.q_device, wires=[k, k+1])\n",
    "            return(self.measure(self.q_device))\n",
    "            \n",
    "    def __init__(self,\n",
    "                 embed_dim: int,\n",
    "                 num_heads: int,\n",
    "                 dropout=0.1,\n",
    "                 mask=None,\n",
    "                 use_bias=False,\n",
    "                 n_qubits: int = 4,\n",
    "                 n_qlayers: int = 1,\n",
    "                 q_device=\"default.qubit\"):\n",
    "        super(MultiHeadAttentionQuantum, self).__init__(embed_dim, num_heads, dropout=dropout, mask=mask, use_bias=use_bias)\n",
    "        \n",
    "        # todo: add intermediate layer to \"dress\" quantum circuit\n",
    "        assert n_qubits == embed_dim, \"Number of qubits ({n_qubits}) does not match embedding dim ({embed_dim})\"\n",
    "        self.n_qubits = n_qubits\n",
    "        self.n_qlayers = n_qlayers\n",
    "        self.dev = tq.QuantumDevice(self.n_qubits)\n",
    "        self.q_layer = self.QLayer()\n",
    "        self.measure = tq.MeasureAll(tq.PauliZ)\n",
    "    def forward(self, x, mask=None):\n",
    "        batch_size, seq_len, embed_dim = x.size()\n",
    "        assert embed_dim == self.embed_dim, f\"Input embedding ({embed_dim}) does not match layer embedding size ({self.embed_dim})\"\n",
    "\n",
    "        K = [self.q_layer(x[:, t, :].clone(),self.dev) for t in range(seq_len)]\n",
    "        Q = [self.q_layer(x[:, t, :].clone(),self.dev) for t in range(seq_len)]\n",
    "        V = [self.q_layer(x[:, t, :].clone(),self.dev) for t in range(seq_len)]\n",
    "        K = torch.Tensor(pad_sequence(K))\n",
    "        Q = torch.Tensor(pad_sequence(Q))\n",
    "        V = torch.Tensor(pad_sequence(V))\n",
    "        x = self.downstream(Q, K, V, batch_size, mask)\n",
    "        output = [self.q_layer(x[:, t, :],self.dev) for t in range(seq_len)]\n",
    "        output = torch.Tensor(pad_sequence(output)).clone()\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75cd180a",
   "metadata": {},
   "source": [
    "## Build FeedForwrad Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "59e73aff",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForwardBase(nn.Module):\n",
    "    def __init__(self, embed_dim, ffn_dim, dropout=0.1):\n",
    "        super(FeedForwardBase, self).__init__()\n",
    "        self.linear_1 = nn.Linear(embed_dim, ffn_dim)\n",
    "        self.linear_2 = nn.Linear(ffn_dim, embed_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    #def forward(self, x):\n",
    "     #   raise NotImplementedError(\"Base class does not implement forward function\")\n",
    "\n",
    "\n",
    "class FeedForwardClassical(FeedForwardBase):\n",
    "    def __init__(self, embed_dim, ffn_dim, dropout=0.1):\n",
    "        super(FeedForwardClassical, self).__init__(embed_dim, ffn_dim, dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.linear_1(x))\n",
    "        x = self.dropout(x)\n",
    "        x = self.linear_2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a8a2cb8",
   "metadata": {},
   "source": [
    "Here, similar to what we have done for the multi-head attention layer, we build the FeedForward quantum layer.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b6da3ec6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForwardQuantum(FeedForwardBase):\n",
    "    class QL(tq.QuantumModule):\n",
    "        def __init__(self):\n",
    "            super().__init__()    \n",
    "            self.n_wires = 8\n",
    "            self.encoder = tq.GeneralEncoder(\n",
    "        [   {'input_idx': [0], 'func': 'rx', 'wires': [0]},\n",
    "            {'input_idx': [1], 'func': 'rx', 'wires': [1]},\n",
    "            {'input_idx': [2], 'func': 'rx', 'wires': [2]},\n",
    "            {'input_idx': [3], 'func': 'rx', 'wires': [2]},\n",
    "            {'input_idx': [4], 'func': 'rx', 'wires': [4]},\n",
    "            {'input_idx': [5], 'func': 'rx', 'wires': [5]},\n",
    "             {'input_idx': [6], 'func': 'rx', 'wires': [6]},\n",
    "          {'input_idx': [7], 'func': 'rx', 'wires': [7]},\n",
    "        ])\n",
    "            self.rx0 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx1 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx2 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx3 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx4 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx5 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx6 = tq.RX(has_params=True, trainable=True)\n",
    "            self.rx7 = tq.RX(has_params=True, trainable=True)\n",
    "            self.measure = tq.MeasureAll(tq.PauliZ)\n",
    "        def forward (self, x, q_device: tq.QuantumDevice):\n",
    "            self.q_device = q_device\n",
    "            self.encoder(self.q_device, x)\n",
    "            self.rx0(self.q_device, wires=0)\n",
    "            self.rx1(self.q_device, wires=1)\n",
    "            self.rx2(self.q_device, wires=2)\n",
    "            self.rx3(self.q_device, wires=3)\n",
    "            self.rx4(self.q_device, wires=4)\n",
    "            self.rx5(self.q_device, wires=5)\n",
    "            self.rx6(self.q_device, wires=6)\n",
    "            self.rx7(self.q_device, wires=7)\n",
    "            for k in range(self.n_wires):\n",
    "                if k==self.n_wires-1:\n",
    "                    tqf.cnot(self.q_device, wires=[k, 0]) \n",
    "                else:\n",
    "                    tqf.cnot(self.q_device, wires=[k, k+1])\n",
    "            return(self.measure(self.q_device))\n",
    "                    \n",
    "    def __init__(self, embed_dim, n_qubits, n_qlayers=1, dropout=0.1, q_device=\"default.qubit\"):\n",
    "        \n",
    "        super(FeedForwardQuantum, self).__init__(embed_dim, ffn_dim=n_qubits, dropout=dropout)\n",
    "        \n",
    "        assert n_qubits == embed_dim, \"Number of qubits ({n_qubits}) does not match embedding dim ({embed_dim})\"\n",
    "        self.n_qubits = n_qubits\n",
    "        self.n_qlayers = n_qlayers\n",
    "        self.dev = tq.QuantumDevice(self.n_qubits)\n",
    "        self.q_l = self.QL()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        batch_size, seq_len, _ = x.size()\n",
    "        x = self.linear_1(x)\n",
    "        o = [self.q_l(x[:, t, :].detach(),self.dev) for t in range(seq_len)]\n",
    "        x = torch.Tensor(pad_sequence(o))\n",
    "        x = self.linear_2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f24c5f7",
   "metadata": {},
   "source": [
    "## Build the whole hybrid transformer model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a5791664",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerBlockBase(nn.Module):\n",
    "    def __init__(self,\n",
    "                 embed_dim: int,\n",
    "                 num_head: int,\n",
    "                 ff_dim: int,\n",
    "                 n_qubits_transformer: int = 0,\n",
    "                 n_qubits_ffn: int = 0,\n",
    "                 n_qlayers: int = 1,\n",
    "                 dropout: float = 0.1,\n",
    "                 mask=None):\n",
    "        super(TransformerBlockBase, self).__init__()\n",
    "        self.attn = None\n",
    "        self.ffn = None\n",
    "        self.dropout1 = nn.Dropout(dropout)\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "        self.norm1 = nn.LayerNorm(embed_dim)\n",
    "        self.norm2 = nn.LayerNorm(embed_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        attn_output = self.attn(x)\n",
    "        x = self.norm1(attn_output + x)\n",
    "        x = self.dropout1(x)\n",
    "\n",
    "        ff_output = self.ffn(x)\n",
    "        x = self.norm2(ff_output + x)\n",
    "        x = self.dropout2(x)\n",
    "\n",
    "        return x\n",
    "class TransformerBlockClassical(TransformerBlockBase):\n",
    "    def __init__(self,\n",
    "                 embed_dim: int,\n",
    "                 num_heads: int,\n",
    "                 ff_dim: int,\n",
    "                 dropout: float = 0.1,\n",
    "                 mask=None):\n",
    "        super(TransformerBlockClassical, self).__init__(embed_dim, num_heads, ff_dim, dropout, mask)\n",
    "        self.attn = MultiHeadAttentionClassical(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout, mask=mask)\n",
    "        self.ffn = FeedForwardClassical(embed_dim, ff_dim)\n",
    "\n",
    "class TransformerBlockQuantum(TransformerBlockBase):\n",
    "    def __init__(self,\n",
    "                 embed_dim: int,\n",
    "                 num_heads: int,\n",
    "                 ffn_dim: int,\n",
    "                 n_qubits_transformer: int = 0,\n",
    "                 n_qubits_ffn: int = 0,\n",
    "                 n_qlayers: int = 1,\n",
    "                 dropout: float = 0.1,\n",
    "                 mask=None,\n",
    "                 q_device='default.qubit'):\n",
    "        super(TransformerBlockQuantum, self).__init__(embed_dim, num_heads, ffn_dim, dropout, mask)\n",
    "        \n",
    "        self.n_qubits_transformer = n_qubits_transformer\n",
    "        self.n_qubits_ffn = n_qubits_ffn\n",
    "        self.n_qlayers = n_qlayers\n",
    "\n",
    "        self.attn = MultiHeadAttentionQuantum(embed_dim,\n",
    "                                              num_heads,\n",
    "                                              n_qubits=n_qubits_transformer,\n",
    "                                              n_qlayers=n_qlayers,\n",
    "                                              dropout=dropout,\n",
    "                                              mask=mask,\n",
    "                                              q_device=q_device)\n",
    "        if n_qubits_ffn > 0:\n",
    "            self.ffn = FeedForwardQuantum(embed_dim, n_qubits_ffn, n_qlayers, q_device=q_device)\n",
    "        else:\n",
    "            self.ffn = FeedForwardClassical(embed_dim, ffn_dim)\n",
    "\n",
    "\n",
    "class PositionalEncoder(nn.Module):\n",
    "    def __init__(self, embed_dim, max_seq_len=512):\n",
    "        super().__init__()\n",
    "        self.embed_dim = embed_dim\n",
    "        \n",
    "        # create constant 'pe' matrix with values dependant on pos and i\n",
    "        pe = torch.zeros(max_seq_len, embed_dim)\n",
    "        for pos in range(max_seq_len):\n",
    "            for i in range(0, embed_dim, 2):\n",
    "                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i)/embed_dim)))\n",
    "                pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1))/embed_dim)))\n",
    "        pe = pe.unsqueeze(0)\n",
    "        self.register_buffer('pe', pe)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # make embeddings relatively larger\n",
    "        x = x * math.sqrt(self.embed_dim)\n",
    "        #add constant to embedding\n",
    "        seq_len = x.size(1)\n",
    "        x = x + torch.autograd.Variable(self.pe[:,:seq_len], requires_grad=False)  # .cuda()\n",
    "        return x\n",
    "\n",
    "\n",
    "class TextClassifier(nn.Module):\n",
    "    def __init__(self,\n",
    "                 embed_dim: int,\n",
    "                 num_heads: int,\n",
    "                 num_blocks: int,\n",
    "                 num_classes: int,\n",
    "                 vocab_size: int,\n",
    "                 ffn_dim: int = 32,\n",
    "                 n_qubits_transformer: int = 0,\n",
    "                 n_qubits_ffn: int = 0,\n",
    "                 n_qlayers: int = 1,\n",
    "                 dropout=0.1,\n",
    "                 q_device=\"device.qubit\"):\n",
    "        super(TextClassifier, self).__init__()\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.num_blocks = num_blocks\n",
    "        self.num_classes = num_classes\n",
    "        self.vocab_size = vocab_size\n",
    "\n",
    "        self.token_embedding = nn.Embedding(vocab_size, embed_dim)\n",
    "        self.pos_embedding = PositionalEncoder(embed_dim)\n",
    "\n",
    "        print(f\"++ There will be {num_blocks} transformer blocks\")\n",
    "\n",
    "        if n_qubits_transformer > 0:\n",
    "            print(f\"++ Transformer will use {n_qubits_transformer} qubits and {n_qlayers} q layers\")\n",
    "            if n_qubits_ffn > 0:\n",
    "                print(f\"The feed-forward head will use {n_qubits_ffn} qubits\")\n",
    "            else:\n",
    "                print(f\"The feed-forward head will be classical\")\n",
    "            print(f\"Using quantum device {q_device}\")\n",
    "\n",
    "            transformer_blocks = [\n",
    "                TransformerBlockQuantum(embed_dim, num_heads, ffn_dim,\n",
    "                                        n_qubits_transformer=n_qubits_transformer,\n",
    "                                        n_qubits_ffn=n_qubits_ffn,\n",
    "                                        n_qlayers=n_qlayers, \n",
    "                                        q_device=q_device) for _ in range(num_blocks)\n",
    "            ]\n",
    "        else:\n",
    "            transformer_blocks = [\n",
    "                TransformerBlockClassical(embed_dim, num_heads, ffn_dim) for _ in range(num_blocks)\n",
    "            ]\n",
    "\n",
    "        self.transformers = nn.Sequential(*transformer_blocks)\n",
    "        if self.num_classes > 2:\n",
    "            self.class_logits = nn.Linear(embed_dim, num_classes)\n",
    "        else:\n",
    "            self.class_logits = nn.Linear(embed_dim, 1)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        tokens = self.token_embedding(x)\n",
    "        # batch_size, seq_len, embed_dim = x.size()\n",
    "        x = self.pos_embedding(tokens)\n",
    "        x = self.transformers(x)\n",
    "        x = x.mean(dim=1)  # global average pooling, works in 1D\n",
    "        x = self.dropout(x)\n",
    "        # x = self.class_logits(x)\n",
    "        # return F.log_softmax(x, dim=1)\n",
    "        return self.class_logits(x)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "edbaae4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "23d20ddb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def binary_accuracy(preds, y):\n",
    "    \"\"\"\n",
    "    Returns accuracy per batch, i.e. if you get 8/10 right, this returns 0.8, NOT 8\n",
    "    \"\"\"\n",
    "\n",
    "    #round predictions to the closest integer\n",
    "    rounded_preds = torch.round(torch.sigmoid(preds))\n",
    "    correct = (rounded_preds == y).float() #convert into float for division \n",
    "    acc = correct.sum() / len(correct)\n",
    "    return acc\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "831f5fe5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion):\n",
    "    epoch_loss = 0\n",
    "    epoch_acc = 0\n",
    "    \n",
    "    model.train()\n",
    "    for batch in iterator:\n",
    "        torch.autograd.set_detect_anomaly(True)\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        inputs = torch.LongTensor(batch.text[0])\n",
    "        if inputs.size(1) > MAX_SEQ_LEN:\n",
    "            inputs = inputs[:, :MAX_SEQ_LEN]\n",
    "        predictions = model(inputs).squeeze(1)\n",
    "        \n",
    "        label = batch.label - 1\n",
    "        #label = label.unsqueeze(1)\n",
    "        loss = criterion(predictions, label)\n",
    "        #loss = F.nll_loss(predictions, label)\n",
    "        acc = binary_accuracy(predictions, label)\n",
    "        \n",
    "        loss.backward(retain_graph=True)\n",
    "        optimizer.step()\n",
    "        \n",
    "        epoch_loss += loss.item()\n",
    "        epoch_acc += acc.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator), epoch_acc / len(iterator)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d39478e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, iterator, criterion):\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    epoch_acc = 0\n",
    "    \n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in iterator:\n",
    "            inputs = torch.LongTensor(batch.text[0])\n",
    "            if inputs.size(1) > MAX_SEQ_LEN:\n",
    "                inputs = inputs[:, :MAX_SEQ_LEN]\n",
    "            predictions = model(inputs).squeeze(1)\n",
    "            \n",
    "            label = batch.label - 1\n",
    "            #label = label.unsqueeze(1)\n",
    "            loss = criterion(predictions, label)\n",
    "            #loss = F.nll_loss(predictions, label)\n",
    "            acc = binary_accuracy(predictions, label)\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "            epoch_acc += acc.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator), epoch_acc / len(iterator)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "596fe2dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def epoch_time(start_time, end_time):\n",
    "    elapsed_time = end_time - start_time\n",
    "    elapsed_mins = int(elapsed_time / 60)\n",
    "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "    return elapsed_mins, elapsed_secs\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b50eeb74",
   "metadata": {},
   "source": [
    "## Training Classical transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "4ebd645a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training examples: 25000\n",
      "Testing examples:  25000\n",
      "++ There will be 1 transformer blocks\n",
      "The model has 160,441 trainable parameters\n",
      "Epoch 1/5\n",
      "Epoch: 01 | Epoch Time: 0m 28s\n",
      "\tTrain Loss: -5.346 | Train Acc: 49.22%\n",
      "\t Val. Loss: -13.767 |  Val. Acc: 50.04%\n",
      "Epoch 2/5\n",
      "Epoch: 02 | Epoch Time: 0m 31s\n",
      "\tTrain Loss: -28.989 | Train Acc: 50.00%\n",
      "\t Val. Loss: -47.243 |  Val. Acc: 50.04%\n",
      "Epoch 3/5\n",
      "Epoch: 03 | Epoch Time: 0m 28s\n",
      "\tTrain Loss: -70.884 | Train Acc: 50.01%\n",
      "\t Val. Loss: -97.031 |  Val. Acc: 50.04%\n",
      "Epoch 4/5\n",
      "Epoch: 04 | Epoch Time: 0m 28s\n",
      "\tTrain Loss: -127.913 | Train Acc: 50.01%\n",
      "\t Val. Loss: -160.671 |  Val. Acc: 50.04%\n",
      "Epoch 5/5\n",
      "Epoch: 05 | Epoch Time: 0m 28s\n",
      "\tTrain Loss: -197.544 | Train Acc: 50.01%\n",
      "\t Val. Loss: -236.139 |  Val. Acc: 50.04%\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "\n",
    "    parser = argparse.ArgumentParser() \n",
    "    parser.add_argument('-D', '--q_device', default='local', type=str)\n",
    "    parser.add_argument('-B', '--batch_size', default=32, type=int)\n",
    "    parser.add_argument('-E', '--n_epochs', default=5, type=int)\n",
    "    parser.add_argument('-C', '--n_classes', default=2, type=int)\n",
    "    parser.add_argument('-l', '--lr', default=0.001, type=float)\n",
    "    parser.add_argument('-v', '--vocab_size', default=20000, type=int)\n",
    "    parser.add_argument('-e', '--embed_dim', default=8, type=int)\n",
    "    parser.add_argument('-s', '--max_seq_len', default=64, type=int)\n",
    "    parser.add_argument('-f', '--ffn_dim', default=8, type=int)\n",
    "    parser.add_argument('-t', '--n_transformer_blocks', default=1, type=int)\n",
    "    parser.add_argument('-H', '--n_heads', default=2, type=int)\n",
    "    parser.add_argument('-q', '--n_qubits_transformer', default=0, type=int)\n",
    "    parser.add_argument('-Q', '--n_qubits_ffn', default=0, type=int)\n",
    "    parser.add_argument('-L', '--n_qlayers', default=1, type=int)\n",
    "    parser.add_argument('-d', '--dropout_rate', default=0.1, type=float)\n",
    "    args = parser.parse_args(args=[])\n",
    "\n",
    "    MAX_SEQ_LEN = args.max_seq_len\n",
    "\n",
    "    TEXT = data.Field(lower=True, include_lengths=True, batch_first=True)\n",
    "    #LABEL = data.Field(sequential=False)\n",
    "    LABEL = data.LabelField(dtype=torch.float)\n",
    "    train_data, test_data = datasets.IMDB.splits(TEXT, LABEL)\n",
    "    print(f'Training examples: {len(train_data)}')\n",
    "    print(f'Testing examples:  {len(test_data)}')\n",
    "\n",
    "    TEXT.build_vocab(train_data, max_size=args.vocab_size - 2)  # exclude <UNK> and <PAD>\n",
    "    LABEL.build_vocab(train_data)\n",
    "\n",
    "    train_iter, test_iter = data.BucketIterator.splits((train_data, test_data), batch_size=args.batch_size)\n",
    "    \n",
    "    model = TextClassifier(embed_dim=args.embed_dim,\n",
    "                           num_heads=args.n_heads,\n",
    "                           num_blocks=args.n_transformer_blocks,\n",
    "                           num_classes=args.n_classes,\n",
    "                           vocab_size=args.vocab_size,\n",
    "                           ffn_dim=args.ffn_dim,\n",
    "                           n_qubits_transformer=args.n_qubits_transformer,\n",
    "                           n_qubits_ffn=args.n_qubits_ffn,\n",
    "                           n_qlayers=args.n_qlayers,\n",
    "                           dropout=args.dropout_rate,\n",
    "                           q_device=args.q_device)\n",
    "    print(f'The model has {count_parameters(model):,} trainable parameters')\n",
    "\n",
    "    optimizer = torch.optim.Adam(lr=args.lr, params=model.parameters())\n",
    "    if args.n_classes < 3:\n",
    "        criterion = torch.nn.BCEWithLogitsLoss()  # logits -> sigmoid -> loss\n",
    "    else:\n",
    "        criterion = torch.nn.CrossEntropyLoss()  # logits -> log_softmax -> NLLloss\n",
    "\n",
    "    # training loop\n",
    "    best_valid_loss = float('inf')\n",
    "    for iepoch in range(args.n_epochs):\n",
    "        start_time = time.time()\n",
    "\n",
    "        print(f\"Epoch {iepoch+1}/{args.n_epochs}\")\n",
    "\n",
    "        train_loss, train_acc = train(model, train_iter, optimizer, criterion)\n",
    "        valid_loss, valid_acc = evaluate(model, test_iter, criterion)\n",
    "\n",
    "        end_time = time.time()\n",
    "\n",
    "        epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "        \n",
    "        if valid_loss < best_valid_loss:\n",
    "            best_valid_loss = valid_loss\n",
    "            torch.save(model.state_dict(), 'model.pt')\n",
    "        \n",
    "        print(f'Epoch: {iepoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')\n",
    "        print(f'\\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc*100:.2f}%')\n",
    "        print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. Acc: {valid_acc*100:.2f}%')\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "211dd509",
   "metadata": {},
   "source": [
    "## Training Quantum transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13683de2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training examples: 25000\n",
      "Testing examples:  25000\n",
      "++ There will be 1 transformer blocks\n",
      "++ Transformer will use 8 qubits and 1 q layers\n",
      "The feed-forward head will use 8 qubits\n",
      "Using quantum device default.qubit\n",
      "The model has 400,201 trainable parameters\n",
      "Epoch 1/1\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "\n",
    "    parser = argparse.ArgumentParser() \n",
    "    parser.add_argument('-D', '--q_device', default='default.qubit', type=str)\n",
    "    parser.add_argument('-B', '--batch_size', default=32, type=int)\n",
    "    parser.add_argument('-E', '--n_epochs', default=1, type=int)\n",
    "    parser.add_argument('-C', '--n_classes', default=2, type=int)\n",
    "    parser.add_argument('-l', '--lr', default=0.001, type=float)\n",
    "    parser.add_argument('-v', '--vocab_size', default=50000, type=int)\n",
    "    parser.add_argument('-e', '--embed_dim', default=8, type=int)\n",
    "    parser.add_argument('-s', '--max_seq_len', default=64, type=int)\n",
    "    parser.add_argument('-f', '--ffn_dim', default=8, type=int)\n",
    "    parser.add_argument('-t', '--n_transformer_blocks', default=1, type=int)\n",
    "    parser.add_argument('-H', '--n_heads', default=2, type=int)\n",
    "    parser.add_argument('-q', '--n_qubits_transformer', default=8, type=int)\n",
    "    parser.add_argument('-Q', '--n_qubits_ffn', default=8, type=int)\n",
    "    parser.add_argument('-L', '--n_qlayers', default=1, type=int)\n",
    "    parser.add_argument('-d', '--dropout_rate', default=0.1, type=float)\n",
    "    args = parser.parse_args(args=[])\n",
    "\n",
    "    MAX_SEQ_LEN = args.max_seq_len\n",
    "\n",
    "    TEXT = data.Field(lower=True, include_lengths=True, batch_first=True)\n",
    "    #LABEL = data.Field(sequential=False)\n",
    "    LABEL = data.LabelField(dtype=torch.float)\n",
    "    train_data, test_data = datasets.IMDB.splits(TEXT, LABEL)\n",
    "    print(f'Training examples: {len(train_data)}')\n",
    "    print(f'Testing examples:  {len(test_data)}')\n",
    "\n",
    "    TEXT.build_vocab(train_data, max_size=args.vocab_size - 2)  # exclude <UNK> and <PAD>\n",
    "    LABEL.build_vocab(train_data)\n",
    "\n",
    "    train_iter, test_iter = data.BucketIterator.splits((train_data, test_data), batch_size=args.batch_size)\n",
    "    \n",
    "    model = TextClassifier(embed_dim=args.embed_dim,\n",
    "                           num_heads=args.n_heads,\n",
    "                           num_blocks=args.n_transformer_blocks,\n",
    "                           num_classes=args.n_classes,\n",
    "                           vocab_size=args.vocab_size,\n",
    "                           ffn_dim=args.ffn_dim,\n",
    "                           n_qubits_transformer=args.n_qubits_transformer,\n",
    "                           n_qubits_ffn=args.n_qubits_ffn,\n",
    "                           n_qlayers=args.n_qlayers,\n",
    "                           dropout=args.dropout_rate,\n",
    "                           q_device=args.q_device)\n",
    "    print(f'The model has {count_parameters(model):,} trainable parameters')\n",
    "\n",
    "    optimizer = torch.optim.Adam(lr=args.lr, params=model.parameters())\n",
    "    if args.n_classes < 3:\n",
    "        criterion = torch.nn.BCEWithLogitsLoss()  # logits -> sigmoid -> loss\n",
    "    else:\n",
    "        criterion = torch.nn.CrossEntropyLoss()  # logits -> log_softmax -> NLLloss\n",
    "\n",
    "    # training loop\n",
    "    best_valid_loss = float('inf')\n",
    "    for iepoch in range(args.n_epochs):\n",
    "        start_time = time.time()\n",
    "\n",
    "        print(f\"Epoch {iepoch+1}/{args.n_epochs}\")\n",
    "\n",
    "        train_loss, train_acc = train(model, train_iter, optimizer, criterion)\n",
    "        valid_loss, valid_acc = evaluate(model, test_iter, criterion)\n",
    "\n",
    "        end_time = time.time()\n",
    "\n",
    "        epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "        \n",
    "        if valid_loss < best_valid_loss:\n",
    "            best_valid_loss = valid_loss\n",
    "            torch.save(model.state_dict(), 'model.pt')\n",
    "        \n",
    "        print(f'Epoch: {iepoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')\n",
    "        print(f'\\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc*100:.2f}%')\n",
    "        print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. Acc: {valid_acc*100:.2f}%')\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c44446be",
   "metadata": {},
   "source": [
    "## Refrences"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6dfef96",
   "metadata": {},
   "source": [
    "[1] Riccardo Di Sipio, Jia-Hong Huang, Samuel Yen-\n",
    "Chi Chen, Stefano Mangini, and Marcel Worring.\n",
    "The dawn of quantum natural language processing.\n",
    "In ICASSP 2022-2022 IEEE International Confer-\n",
    "ence on Acoustics, Speech and Signal Processing\n",
    "(ICASSP), pages 8612–8616. IEEE, 2022."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "529c8cfd",
   "metadata": {},
   "source": [
    "[2] Ashish Vaswani, Noam Shazeer, Niki Parmar,\n",
    "Jakob Uszkoreit, Llion Jones, Aidan N Gomez,\n",
    "Łukasz Kaiser, and Illia Polosukhin. Attention is all\n",
    "you need. In Advances in neural information processing systems, pages 5998–6008, 2017."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
